package cc.shacocloud.luanniao.core.service.text.splitter

import cc.shacocloud.luanniao.core.service.ai.AiModel
import org.slf4j.Logger
import org.slf4j.LoggerFactory

/**
 * 递归地尝试按不同的字符进行拆分，以找到一个有效的字符
 *
 * @author 思追(shaco)
 */
class RecursiveCharacterTextSplitterServiceImpl(

    /**
     *  分段最大长度
     */
    val maxTokens: Int = 500,

    /**
     * 分段重叠长度
     */
    val chunkOverlap: Int = 50,

    /**
     * 使用的 ai 模型
     */
    val aiModel: AiModel,
) : TextSplitterService {

    companion object {
        private val log: Logger = LoggerFactory.getLogger(RecursiveCharacterTextSplitterServiceImpl::class.java)
    }

    /**
     * 分割文本
     *
     * @param text 文本
     * @param separators 分隔符
     */
    override fun splitText(
        text: String,
        separators: List<String>,
        keepSeparator: Boolean,
    ): List<String> {
        require(separators.isNotEmpty()) { "分隔符必须不为空" }


        // 匹配合适的分隔符
        var separator = separators[0]
        var newSeparators: List<String> = listOf()
        for ((i, s) in separators.withIndex()) {
            if (s == "") {
                separator = s
                break
            } else if (text.matches(separator.toRegex())) {
                separator = s
                newSeparators = separators.subList(i + 1, separators.size)
            }
        }

        val chunks = regexSplitText(text, separator, keepSeparator)
        val result = mutableListOf<String>()

        // 递归分割较长的文本
        separator = if (keepSeparator) "" else separator
        val tmpSplits = mutableListOf<String>()
        for (chunk in chunks) {
            val tokenLen = aiModel.textToTokens(chunk).size
            if (tokenLen < maxTokens) {
                tmpSplits.add(chunk)
            } else {

                if (tmpSplits.isNotEmpty()) {
                    val mergedText = mergeSplits(tmpSplits, separator)
                    result.addAll(mergedText)
                    tmpSplits.clear()
                }
                if (newSeparators.isEmpty()) {
                    result.add(chunk)
                } else {
                    val otherInfo = splitText(chunk, newSeparators, keepSeparator)
                    result.addAll(otherInfo)
                }
            }
        }

        if (tmpSplits.isNotEmpty()) {
            val mergedText = mergeSplits(tmpSplits, separator)
            result.addAll(mergedText)
        }

        return result
    }

    fun mergeSplits(splits: List<String>, separator: String): List<String> {
        val separatorLen = aiModel.textToTokens(separator).size

        val docs = mutableListOf<String>()
        var currentDoc = mutableListOf<String>()
        var total = 0

        for (str in splits) {
            val len = aiModel.textToTokens(str).size
            if (total + len + (if (currentDoc.isNotEmpty()) separatorLen else 0) > maxTokens) {
                if (total > maxTokens) {
                    if (log.isWarnEnabled) {
                        log.warn("创建了一个大小为 {} 的块，该块比指定的 {} 长", total, maxTokens)
                    }
                }

                if (currentDoc.isNotEmpty()) {
                    val doc = joinDocs(currentDoc, separator)
                    if (doc != null) {
                        docs.add(doc)
                    }

                    while (total > chunkOverlap || (total + len + (if (currentDoc.size > 1) separatorLen else 0) > maxTokens && total > 0)) {
                        total -= aiModel.textToTokens(currentDoc.first()).size + (if (currentDoc.size > 1) separatorLen else 0)
                        currentDoc = currentDoc.drop(1).toMutableList()
                    }
                }
            }
            currentDoc.add(str)
            total += len + (if (currentDoc.size > 1) separatorLen else 0)
        }

        val doc = joinDocs(currentDoc, separator)
        if (doc != null) {
            docs.add(doc)
        }
        return docs
    }

    fun joinDocs(docs: List<String>, separator: String): String? {
        val text = docs.joinToString(separator = separator).trim()
        return text.ifEmpty { null }
    }

    /**
     * 使用正则拆分文本
     */
    fun regexSplitText(
        text: String,
        separator: String,
        keepSeparator: Boolean,
    ): List<String> {

        val chunks = if (keepSeparator) {
            mutableListOf<String>().apply {

                // 模式中的括号将分隔符保留在结果中
                val splits = "(${Regex.escape(separator)})".toRegex().split(text)

                add(splits.first())

                // 包含 splits 中每两个连续元素相加的结果
                splits.windowed(2, 2, false)
                    .forEach { add(it.first() + it.last()) }

                // 如果 splits 的长度是偶数，将最后一个元素添加到列表
                if (splits.size % 2 == 0) {
                    add(splits.last())
                }
            }
        } else {
            separator.toRegex().split(text)
        }

        return chunks.filter { it.isNotBlank() }
    }
}